from dataclasses import dataclass, field, asdict
from typing import Any, Dict, List, Optional, Union

import torch

from continuallearning.models.pefts.common.config.hook_manager_config import (
    ModuleDiscoveryConfig,
)


@dataclass
class BackboneFeatures:
    """Container for backbone features collected during forward pass."""

    features: Dict[str, torch.Tensor] = field(default_factory=dict)
    layer_names: List[str] = field(default_factory=list)
    batch_size: Optional[int] = None
    device: Optional[torch.device] = None

    def add_feature(self, layer_name: str, feature: torch.Tensor) -> None:
        """Add a feature tensor for a specific layer."""
        self.features[layer_name] = feature
        if layer_name not in self.layer_names:
            self.layer_names.append(layer_name)

        if self.batch_size is None:
            self.batch_size = feature.shape[0]
        if self.device is None:
            self.device = feature.device

    def get_feature(self, layer_name: str) -> Optional[torch.Tensor]:
        """Get feature tensor for a specific layer."""
        return self.features.get(layer_name)

    def clear(self) -> None:
        """Clear all stored features."""
        self.features.clear()
        self.layer_names.clear()
        self.batch_size = None
        self.device = None


@dataclass
class TaskContext:
    """Container for task-related context information."""

    task_ids: Optional[List[int]] = None
    backbone_features: Optional[BackboneFeatures] = None
    metadata: Dict[str, Any] = field(default_factory=dict)

    def add_metadata(self, key: str, value: Any) -> None:
        """Add metadata to the context."""
        self.metadata[key] = value

    def get_metadata(self, key: str, default: Any = None) -> Any:
        """Get metadata from the context."""
        return self.metadata.get(key, default)

    def clear(self) -> None:
        """Clear all context information."""
        self.task_ids = None
        if self.backbone_features:
            self.backbone_features.clear()
            self.backbone_features = None
        self.metadata.clear()


@dataclass
class ContextManagerConfig:
    """Complete configuration for context manager."""

    # Module discovery configuration
    discovery: ModuleDiscoveryConfig = field(default_factory=ModuleDiscoveryConfig)
    extract_input: bool = False
    extract_output: bool = True

    def __post_init__(self):
        """Process nested configurations."""
        # Convert dict to ModuleDiscoveryConfig if needed
        if isinstance(self.discovery, dict):
            self.discovery = ModuleDiscoveryConfig(**self.discovery)

    def validate(self) -> None:
        """Validate hook manager configuration."""
        # Validate nested configurations
        self.discovery.validate()

    def to_dict(self) -> Dict[str, Any]:
        """Convert configuration to a dictionary."""
        # Get dictionary representation of the dataclass
        config_dict = asdict(self)
        return config_dict

    @classmethod
    def from_dict(cls, config_dict: Dict[str, Any]):
        """Create configuration from a dictionary."""
        # Filter to include only fields that are part of this dataclass
        filtered_dict = {
            k: v
            for k, v in config_dict.items()
            if k in [field.name for field in cls.__dataclass_fields__.values()]
        }

        return cls(**filtered_dict)
